from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report
import torch
import torch.nn as nn
import torch.optim as optim
from matplotlib import pyplot as plt
import numpy as np

#加载数据
iris = load_iris()
x_train, x_test, y_train, y_test = train_test_split(iris['data'], iris['target'], test_size = 0.25)
x_train, x_test = torch.tensor(x_train, dtype=torch.float), torch.tensor(x_test, dtype=torch.float)
y_train, y_test = torch.tensor(y_train, dtype=torch.long), torch.tensor(y_test, dtype=torch.long)

#确定网络
net = nn.Sequential(
    nn.Linear(4, 10),
    nn.ReLU(),
    nn.Linear(10, 3),
)

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.04)

#训练
error = []
for epoch in range(150):
    y_pred = net(x_train)
    loss = criterion(y_pred, y_train)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    error.append(loss.item())

#评估
y_pred = net(x_test)
y_pred = torch.argmax(y_pred, dim=1)

print("======== 神经网络 =========")
print(classification_report(y_test, y_pred, target_names=iris['target_names']))

plt.plot(error)
plt.show()



